import numpy as np
import cv2
from collections import Counter


def floodfill_alpha(img_array, bg_colors, tolerance, feather, edge_clarity):

    # 取高度宽度
    height, width = img_array.shape[:2]

    # 根据是否传入取样颜色决定背景色来源
    if bg_colors:
        # 用户已手动取样：直接使用取样颜色作为 flood-fill 种子
        background_colors = [tuple(map(int, c)) for c in bg_colors]
    else:
        # ------------- 自动检测背景颜色（复制自原算法） -------------
        edge_pixels = []
        edge_pixels.extend(img_array[0, :, :3].tolist())          # 上边缘
        edge_pixels.extend(img_array[height - 1, :, :3].tolist()) # 下边缘
        edge_pixels.extend(img_array[:, 0, :3].tolist())          # 左边缘
        edge_pixels.extend(img_array[:, width - 1, :3].tolist())  # 右边缘

        edge_pixels = [tuple(pixel) for pixel in edge_pixels]
        color_counts = Counter(edge_pixels)
        total_edge_pixels = len(edge_pixels)

        threshold_percentage = 1
        background_colors = []
        for color, count in color_counts.items():
            if (count / total_edge_pixels) * 100 >= threshold_percentage:
                background_colors.append(color)
        if not background_colors:
            background_colors = [color_counts.most_common(1)[0][0]]
        # -----------------------------------------------------------

    # 创建掩码等后续处理
    background_mask = np.zeros((height, width), dtype=np.uint8)
    visited = np.zeros((height, width), dtype=bool)
    bg_colors_array = np.array(background_colors)

    def is_background_color(pixel):
        distances = np.abs(bg_colors_array - pixel).sum(axis=1)
        return np.any(distances <= tolerance * 3)

    queue = []
    for x in range(width):
        if is_background_color(img_array[0, x, :3]):
            queue.append((x, 0))
            visited[0, x] = True
            background_mask[0, x] = 255
        if is_background_color(img_array[height - 1, x, :3]):
            queue.append((x, height - 1))
            visited[height - 1, x] = True
            background_mask[height - 1, x] = 255
    for y in range(height):
        if not visited[y, 0] and is_background_color(img_array[y, 0, :3]):
            queue.append((0, y))
            visited[y, 0] = True
            background_mask[y, 0] = 255
        if not visited[y, width - 1] and is_background_color(img_array[y, width - 1, :3]):
            queue.append((width - 1, y))
            visited[y, width - 1] = True
            background_mask[y, width - 1] = 255

    directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
    idx = 0
    while idx < len(queue):
        x, y = queue[idx]
        idx += 1
        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            if 0 <= nx < width and 0 <= ny < height and not visited[ny, nx]:
                if is_background_color(img_array[ny, nx, :3]):
                    queue.append((nx, ny))
                    visited[ny, nx] = True
                    background_mask[ny, nx] = 255

    # 前景掩码、形态学清理
    foreground_mask = cv2.bitwise_not(background_mask)
    kernel = np.ones((3, 3), np.uint8)
    cleaned_mask = cv2.morphologyEx(foreground_mask, cv2.MORPH_OPEN, kernel)
    cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel)

    # 对比度增强
    alpha_val = float(edge_clarity) * 0.2
    enhanced_mask = cv2.convertScaleAbs(cleaned_mask, alpha=alpha_val, beta=0)

    # Canny 边缘
    edges = cv2.Canny(enhanced_mask, 50, 150)
    sure_fg = cv2.erode(enhanced_mask, kernel, iterations=2)
    sure_bg = cv2.dilate(enhanced_mask, kernel, iterations=2)
    refined_fg = cv2.bitwise_or(sure_fg, edges)

    trimap = np.zeros((height, width), dtype=np.uint8)
    trimap[sure_bg == 0] = 0
    trimap[refined_fg > 0] = 255
    uncertain = cv2.bitwise_xor(sure_bg, refined_fg)
    trimap[uncertain > 0] = 128

    # 距离变换平滑过渡
    dist_transform = cv2.distanceTransform(uncertain, cv2.DIST_L2, 3)
    cv2.normalize(dist_transform, dist_transform, 0, 1.0, cv2.NORM_MINMAX)
    alpha_uncertain = (dist_transform * 255).astype(np.uint8)

    final_alpha = np.zeros((height, width), dtype=np.uint8)
    final_alpha[trimap == 255] = 255
    final_alpha[trimap == 0] = 0
    final_alpha[trimap == 128] = alpha_uncertain[trimap == 128]

    # 羽化
    if feather > 0:
        edge_area = cv2.dilate(refined_fg, kernel, iterations=int(feather)) - cv2.erode(refined_fg, kernel, iterations=int(feather))
        edge_alpha = final_alpha.copy()
        blur_size = int(feather * 2) * 2 + 1
        edge_alpha = cv2.GaussianBlur(edge_alpha, (blur_size, blur_size), feather)
        final_alpha[edge_area > 0] = edge_alpha[edge_area > 0]

    # 双边滤波
    final_alpha = cv2.bilateralFilter(final_alpha, 9, 75, 75)

    return final_alpha